"""
https://blog.csdn.net/menghaocheng/article/details/102783705

【TF2.0-CNN】迁移学习（将inceptionV3应用到猫狗分类）
"""

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import os

VER = 'v1.0'
N_EPOCHS = 4
FILE_NAME = os.path.basename(__file__)
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
SAVE_PATH = os.path.join(SAVE_DIR, 'trans_learn_self_model.dat')
LOG_DIR = os.path.join('_log', FILE_NAME, VER)

local_weights_file = '../../../../../large_data/model/inceptionV3/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5'

pre_trained_model = tf.keras.applications.inception_v3.InceptionV3(input_shape=(150, 150, 3),
                                                                   include_top=False,
                                                                   weights=None)
TRAIN_DIR = '../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/train_fast'
VAL_DIR = '../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/validation_fast'

pre_trained_model.load_weights(local_weights_file)
for layer in pre_trained_model.layers:
    layer.trainable = False  # ATTENTION

last_layer = pre_trained_model.get_layer('mixed7')
print('last layer output shape: ', last_layer.output_shape)
last_output = last_layer.output
x = tf.keras.layers.Flatten()(last_output)
x = tf.keras.layers.Dense(1024, activation='relu')(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.Dense(1, activation='sigmoid')(x)

model = tf.keras.Model(pre_trained_model.input, x)

model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=0.0001),
              loss='binary_crossentropy',
              metrics=['acc'])

train_datagen = ImageDataGenerator(rescale=1. / 255.,
                                   rotation_range=40,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True)

train_generator = train_datagen.flow_from_directory(TRAIN_DIR,
                                                    batch_size=20,
                                                    class_mode='binary',
                                                    target_size=(150, 150))

test_datagen = ImageDataGenerator(rescale=1.0 / 255.)
validation_generator = test_datagen.flow_from_directory(VAL_DIR,
                                                        batch_size=20,
                                                        class_mode='binary',
                                                        target_size=(150, 150))

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0)

history = model.fit_generator(
    train_generator,
    validation_data=validation_generator,
    # steps_per_epoch=100,
    steps_per_epoch=None,
    # epochs=20,
    epochs=N_EPOCHS,
    # validation_steps=50,
    validation_steps=None,
    verbose=2,
    callbacks=tensorboard_callback,
)
